# Grid trading strategies
# 
# Author: Mario
# 
# Version: v.20200424
#
import logging
import numpy as np
from collections import defaultdict

logger = logging.getLogger()
logger.setLevel(logging.INFO)

class Account:
	# 账户的作用是:
	# 1) 分别在盘前和盘中记录仓位情况, 同时进行仓位的管理
	# 2) hist_account 历史账户主要记录日前所有仓位的信息，以便在盘前函数中做规划
	# 3) daily_account 是日内账户，主要记录日内所有仓位的信息，每日刷新
	# 账户结构: 
	# 1) 历史账户: dict(int)
	# 2) 日内账户: dict(int)

	def __init__(self, capital, logger, min_, max_, power):
		# 账户本金:
		self.capital = capital
		# 获得历史价格区间:
		self.min_price = min_
		self.max_price = max_
		# 格子的幂次单位:
		self.power = power
		self.unit = 10 ** power
		# 日志记录:
		self.logger = logger
		# 生成历史账户的价格序列:
		assert self.min_price != self.max_price, \
			   "invalid price range: [%.3f, %.3f]" % (self.min_price, self.max_price)
		# 生成历史账户:
		self.hist_account = {price: 0 for price in self._sequence_prices(min_, max_)}
		self.daily_account = defaultdict(int)
		
	def _get_prices(self):
		# :return: 返回历史账户中的价格数组
		return np.array(list(self.hist_account.keys()))	
	
	def update_capital(self, capital_flow):
		self.capital = capital_flow if capital_flow > 0 else 0

	def can_trade(self, volume):
		# 查询本金是否满足交易条件:
		return True if self.capital >= volume else False

	def update_price_range(self, min_, max_):
		self.min_price = min_
		self.max_price = max_
		# 修改账户序列, 当且仅当价格序列不发生冲突时, 冲突存在于
		# 当修改的价格小于原价格序列的最高值且仍有持仓时:
		self.logger.info("[alter price] min: %.3f, max: %.3f" % (min_, max_))
	
	def _record_check(self, record):
		assert record.get("quantity", 0), "invalid quantity: %s" % str(record)
		assert record.get("price", 0), "invalid price: %s" % str(record)
		# 统一价格和股数的小数位:
		record["price"] = round(record["price"], self.power)
		record["quantity"] = int(record["quantity"])
		return record
	
	def hist_update(self, record):
		# 对于历史账户的修改只能通过update函数进行
		# 传入的record结构: dict(price=float, quantity=int)
		self._record_check(record)
		# 每次更新前，先对可能存在的重复价格进行合并:
		price = round(record["price"], self.power)
		quantity = int(record["quantity"])
		self.hist_account[price] += quantity

	def daily_update(self, record):
		self._record_check(record)
		self.daily_account[record["price"]] += record["quantity"]
	
	def daily_cleanup(self):
		# 归零当日账户并添加到历史账户中:
		# :param today: 格式为 "20001010"
		for price, quantity in self.daily_account.items():
			self.hist_update(dict(price=price, quantity=quantity))
		# 日内账户归零:
		self.daily_account = defaultdict(int)
		
	def _sequence_prices(self, low, high):
		# :return: 返回一个价格序列(float)
		return np.arange(int(low * 10 ** self.power),
				  int(high * 10 ** self.power),
				  1) / 10 ** self.power

	def _query(self, price):
		# 历史账户的查询操作:
		# 1) 价格落在区间中时，向下求和;
		# 2) 价格超出区间时: 
		#   a) 低于区间给出信号"补仓-merge", 满足条件时购入低价格子;
		#   b) 高于区间时给出价差信号"价差-gap", 如果满足出清条件则清仓.
		# 3) 明确返回的格式: (可以根据策略对购入量实行梯度分布)
		""" 当存在价差时的账户查询返回一个标准的JSON对象
		{0.500: 100,
		 ... ...
		 0.510: 100}
		"""
		prices = np.array(list(self.hist_account.keys()))
		if price > prices.max():
			return {price: int(sel.hist_account.quantity.sum())}
		
		elif price < prices.min():
			return {price: int(self.capital // (price * 10**self.power))} 
			
		else:
			pack = defaultdict(int)
			for pr, share in self.hist_account.items():
				if pr <= price:
					pack[pr] += int(share)
			return pack
			
	def daily_query(self, price):
		price = round(price, self.power)
		return dict(price=price,
					quantity=self.daily_account[price] if price in self.daily_account else 0)
	
	def hist_query(self, price):
		return self._query(price)
		
